import xml.etree.ElementTree as ET
import os

with open('voc_classes.txt') as f:
    classes = f.readlines()
classes = [c.strip() for c in classes]

def get_truthbox(image_id, list_file):
    in_file = open('input/ground-truth-xml/%s.xml'%(image_id), encoding='UTF-8')
    tree=ET.parse(in_file)
    root = tree.getroot()

    for obj in root.iter('object'):
        difficult = obj.find('difficult').text
        cls = obj.find('name').text
        if cls == 'complex_crack':
            cls = 'simple_crack'
        if cls not in classes or int(difficult)==1:
            print("Attention: " + str(cls) + "is not in classes, please make sure of it!!!")
            continue
        xmlbox = obj.find('bndbox')
        b = (int(xmlbox.find('xmin').text), int(xmlbox.find('ymin').text), int(xmlbox.find('xmax').text), int(xmlbox.find('ymax').text)) # left top right bottm
        list_file.write(str(cls) + " " + " ".join([str(a) for a in b]) + '\r')
    print('-----OK-----')


if __name__ == "__main__":
    wd = os.getcwd()
    xml_path = 'input/ground-truth-xml/'
    for i, image_ids in enumerate(os.listdir(xml_path)):
        image_id = image_ids.split('.')[0]
        suffix = image_ids.split('.')[-1]
        print(suffix)
        print(id(suffix), id('xml'))
        if suffix != 'xml':
            continue
        list_file = open('input/ground-truth/%s.txt'%(image_id), 'w')
        get_truthbox(image_id, list_file)
        list_file.close()

